package com.xzy.es.demo.doc;

import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.bulk.BulkItemResponse;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;

/**
 * 文档相关操作
 *
 * @author xzy
 * @date 2021/10/2110:16
 */
@Slf4j
public class DocMain {

    /**
     * 创建文档
     *
     * @param esClient  - ES客户端
     * @param indexName - 索引名称
     * @param docId     - 文档ID
     * @param docObject - 文档
     * @throws IOException -
     */
    public static void insertDoc(RestHighLevelClient esClient, String indexName, String docId, Object docObject) throws IOException {
        IndexRequest indexRequest = new IndexRequest();
        indexRequest
                // 指定文档所属的索引
                .index(indexName)
                // 手动设置文档的ID
                .id(docId)
                // JSON格式的文档
                .source(JSON.toJSONString(docObject), XContentType.JSON);
        IndexResponse indexResponse = esClient.index(indexRequest, RequestOptions.DEFAULT);
        log.info("文档创建结果：\nresult:{}\nversion:{}\n", indexResponse.getResult(), indexResponse.getVersion());
    }

    /**
     * 创建文档（批量）
     *
     * @param esClient   - ES客户端
     * @param indexName  - 索引名称
     * @param docObjects - 文档
     * @return - 文档ID
     * @throws IOException -
     */
    public static List<String> insertDocInBatch(RestHighLevelClient esClient, String indexName, List<Object> docObjects) throws IOException {

        BulkRequest bulkRequest = new BulkRequest();
        for (Object docObject : docObjects) {
            bulkRequest.add(
                    new IndexRequest()
                            .index(indexName)
                            .id(UUID.randomUUID().toString())
                            .source(JSON.toJSONString(docObject), XContentType.JSON)
            );
        }

        BulkResponse bulkResponse = esClient.bulk(bulkRequest, RequestOptions.DEFAULT);

        log.info("批量新增总耗时：{}", bulkResponse.getTook());
        log.info("批量新增结果：");
        List<String> docIds = new ArrayList<>(docObjects.size());
        for (BulkItemResponse item : bulkResponse.getItems()) {
            log.info("index:{},id:{}", item.getIndex(), item.getId());
            docIds.add(item.getId());
        }

        return docIds;
    }

    /**
     * 修改文档的指定域
     *
     * @param esClient  - ES客户端
     * @param indexName - 索引名称
     * @param docId     - 文档ID
     * @param field     - 文档中需要修改的域
     * @param value     - 修改后的值
     * @throws IOException -
     */
    public static void updateDoc(RestHighLevelClient esClient, String indexName, String docId, Object field, Object value) throws IOException {
        UpdateRequest updateRequest = new UpdateRequest();
        updateRequest
                .index(indexName)
                .id(docId)
                .doc(XContentType.JSON, field, value);
        UpdateResponse updateResponse = esClient.update(updateRequest, RequestOptions.DEFAULT);
        log.info("文档修改结果：\nresult:{}\nversion:{}\n", updateResponse.getResult(), updateResponse.getVersion());
    }

    /**
     * 根据ID查询文档
     *
     * @param esClient  - ES客户端
     * @param indexName - 索引名称
     * @param docId     - 文档ID
     * @throws IOException -
     */
    public static void selectDoc(RestHighLevelClient esClient, String indexName, String docId) throws IOException {
        GetRequest getRequest = new GetRequest();
        getRequest
                .index(indexName)
                .id(docId);
        GetResponse getResponse = esClient.get(getRequest, RequestOptions.DEFAULT);
        log.info("查询到的文档信息：\nindex:{}\nid:{}\nsource:{}\n", getResponse.getIndex(), getResponse.getId(), getResponse.getSourceAsString());
    }

    /**
     * 查询指定索引下的所有文档
     *
     * @param esClient  - ES客户端
     * @param indexName - 索引名称
     * @throws IOException -
     */
    public static void findAllDoc(RestHighLevelClient esClient, String indexName, int from, int size) throws IOException {
        SearchRequest searchRequest = new SearchRequest();
        // 设置查询的索引
        searchRequest.indices(indexName);
        // 设置文档查询、分页条件
        searchRequest.source(new SearchSourceBuilder()
                .query(QueryBuilders.matchAllQuery())
                .from(from)
                .size(size)
        );

        // 检索文档
        SearchResponse searchResponse = esClient.search(searchRequest, RequestOptions.DEFAULT);

        // 检索结果
        SearchHits searchHits = searchResponse.getHits();

        log.info("查询总耗时：{}  匹配数据条数：{}", searchResponse.getTook(), searchHits.getTotalHits().value);
        for (SearchHit searchHit : searchHits) {
            log.info(searchHit.getSourceAsString());
        }
    }

    /**
     * 根据ID删除文档
     *
     * @param esClient  - ES客户端
     * @param indexName - 索引名称
     * @param docId     - 文档ID
     * @throws IOException -
     */
    public static void deleteDoc(RestHighLevelClient esClient, String indexName, String docId) throws IOException {
        DeleteRequest deleteRequest = new DeleteRequest();
        deleteRequest
                .index(indexName)
                .id(docId);

        DeleteResponse deleteResponse = esClient.delete(deleteRequest, RequestOptions.DEFAULT);
        log.info("文档删除结果：\n{}\n", deleteResponse);
    }

    /**
     * 根据ID删除文档（批量）
     *
     * @param esClient  - ES客户端
     * @param indexName - 索引名称
     * @param docIds    - 文档ID
     * @throws IOException -
     */
    public static void deleteDocInBatch(RestHighLevelClient esClient, String indexName, List<String> docIds) throws IOException {
        BulkRequest bulkRequest = new BulkRequest();
        for (String docId : docIds) {
            bulkRequest.add(new DeleteRequest().index(indexName).id(docId));
        }

        BulkResponse bulkResponse = esClient.bulk(bulkRequest, RequestOptions.DEFAULT);

        log.info("批量删除总耗时：{}", bulkResponse.getTook());
        log.info("批量删除结果：");
        for (BulkItemResponse item : bulkResponse.getItems()) {
            log.info("index:{},id:{}", item.getIndex(), item.getId());
        }
    }
}
